import json
import time
import re
import os
from openai import OpenAI
from docx import Document

# ========================== 1. 配置区 ==========================
API_KEY = "BearerCubJNNWPCQjooxiEbLGB:IgSlCSEPDRuFXHJlXlMw" 
URL = "https://spark-api-open.xf-yun.com/v2/chat/completions"
INPUT_DOCUMENT_PATH = r"C:/Users/cassi/Desktop/bisai/word-master/豫洛阳-兰郑长干线-CPY-0790-BFGDGS-ZZSYQFGS.docx"
OUTPUT_JSON_PATH = r"C:/Users/cassi/Desktop/bisai/word-master/data_logic_verification.json"

# 数据验证规则
VALIDATION_RULES = {
    'potential_range': (0.85, 1.2),  # 电位测试结果范围（正电压）
    'risk_levels': ['低', '中', '较高', '高'],  # 风险等级
    'numeric_fields': ['可能性', '后果值', '风险值']
}

# ========================== 2. 功能函数区 ==========================
def get_answer(message):
    try:
        client = OpenAI(
            api_key="sk-ma0HIkN3beK0mtThtDzIfRLN5IzpL4agmB9RSbpHk53h7BF5",
            base_url="https://www.chataiapi.com/v1",
        )
        response = client.chat.completions.create(
            model="gemini-2.5-pro",
            messages=message,
            temperature=0.7,
        )
        content = response.choices[0].message.content or ""
        if content.strip():
            print("\n\n*******************模型生成内容如下********************\n")
            print(content)
        return content
    except Exception as e:
        print(f"API请求失败：{str(e)}")
        return ""

def extract_text_from_docx(docx_path):
    """从Word文档中提取文本内容"""
    try:
        doc = Document(docx_path)
        full_text = ""
        
        # 提取段落文本
        for para in doc.paragraphs:
            if para.text.strip():
                full_text += para.text.strip() + "\n"
        
        # 提取表格文本
        for table in doc.tables:
            for row in table.rows:
                for cell in row.cells:
                    if cell.text.strip():
                        full_text += cell.text.strip() + " "
                full_text += "\n"
        
        return full_text
    except Exception as e:
        print(f"Word文档文本提取失败：{str(e)}")
        return ""

def extract_context(text, start, end, window=150):
    """基于句界与换行智能提取上下文，尽量返回完整句子或段落。
    优先以中文句号/问号/叹号及换行符为边界，退化为固定窗口。
    """
    # 左边界：最近的句末或换行
    left_candidates = [
        text.rfind("。", 0, start),
        text.rfind("！", 0, start),
        text.rfind("？", 0, start),
        text.rfind("\n", 0, start)
    ]
    left = max([pos for pos in left_candidates if pos != -1], default=-1)
    if left == -1:
        left = max(0, start - window)
    else:
        left = left + 1  # 从标点后一位开始

    # 右边界：下一处句末或换行
    right_candidates = [
        text.find("。", end),
        text.find("！", end),
        text.find("？", end),
        text.find("\n", end)
    ]
    right_candidates = [pos for pos in right_candidates if pos != -1]
    if right_candidates:
        right = min(right_candidates) + 1
    else:
        right = min(len(text), end + window)

    return text[left:right].strip()
def extract_potential_test_results(text):
    """
    改进版：仅提取真实电位值（单位V或含正负号，范围在[-2, 2]之间）
    兼容格式：
    - '电位：-0.95V'
    - 'K79 +0.989 +0.994 +1.002'
    - '近6月 0.97 0.98 1.00'
    """
    import re
    potential_results = []
    clean_text = re.sub(r'\s+', ' ', text)

    # 提取带有“电位”、“V”、“伏”、“K号”的上下文
    patterns = [
        r'(?:电位|IR降|K\d+)[^0-9\-+]{0,10}?([+\-]?\d+(?:\.\d+)?)\s*[Vv伏]?',
        r'(?:电位测试结果|近\d+月)[：:\s]*([+\-]?\d+(?:\.\d+)?)'
    ]

    for pattern in patterns:
        for m in re.finditer(pattern, clean_text):
            try:
                val = float(m.group(1))
                # 排除异常值（年份、编号）
                if not (-2.0 <= val <= 2.0):
                    continue
                context = clean_text[max(0, m.start()-80):m.end()+80]
                potential_results.append({
                    'value': val,
                    'context': context.strip()
                })
            except:
                continue

    return potential_results


def extract_risk_assessment_data(text):
    """提取风险评价数据（增强版，兼容表格、无冒号格式、失效字段）"""
    clean_text = re.sub(r'\s+', ' ', text)
    risk_data = {
        'possibility_values': [],
        'consequence_values': [],
        'risk_values': [],
        'risk_levels': []
    }

    # 可能性
    for match in re.finditer(r'(?:失效)?可能性[：:\s]*([0-9]*\.?[0-9]+)', clean_text):
        risk_data['possibility_values'].append({
            'value': float(match.group(1)),
            'raw_text': match.group(0),
            'context': clean_text[max(0, match.start()-80):match.end()+80]
        })

    # 后果值
    for match in re.finditer(r'(?:失效)?后果值[：:\s]*([0-9]*\.?[0-9]+)', clean_text):
        risk_data['consequence_values'].append({
            'value': float(match.group(1)),
            'raw_text': match.group(0),
            'context': clean_text[max(0, match.start()-80):match.end()+80]
        })

    # 风险值
    for match in re.finditer(r'风险值[：:\s]*([0-9]*\.?[0-9]+)', clean_text):
        risk_data['risk_values'].append({
            'value': float(match.group(1)),
            'raw_text': match.group(0),
            'context': clean_text[max(0, match.start()-80):match.end()+80]
        })

    # 风险等级
    for match in re.finditer(r'风险等级[：:\s]*([低中较高高]+)', clean_text):
        risk_data['risk_levels'].append({
            'level': match.group(1).strip(),
            'raw_text': match.group(0),
            'context': clean_text[max(0, match.start()-80):match.end()+80]
        })

    return risk_data


def verify_data_logic(potential_results, risk_data):
    """验证数据逻辑正确性"""
    verification_results = {
        'potential_valid': True,
        'risk_values_numeric': True,
        'risk_levels_valid': True,
        'data_logic_score': 5,
        'deduction_points': 0,
        'issues': [],
        'data_summary': {}
    }
    
    # 检查电位测试结果
    valid_potentials = 0
    invalid_potentials = []
    
    for result in potential_results:
        value = result['value']
        if VALIDATION_RULES['potential_range'][0] <= value <= VALIDATION_RULES['potential_range'][1]:
            valid_potentials += 1
        else:
            invalid_potentials.append(result)
    
    if invalid_potentials:
        verification_results['potential_valid'] = False
        verification_results['issues'].append(f"电位测试结果超出范围：{[r['value'] for r in invalid_potentials]}")
        verification_results['data_logic_score'] -= 2
    
    verification_results['data_summary']['potential_results'] = {
        'total': len(potential_results),
        'valid': valid_potentials,
        'invalid': len(invalid_potentials)
    }
    
    # 检查风险评价数值字段
    numeric_checks = {
        'possibility': risk_data['possibility_values'],
        'consequence': risk_data['consequence_values'], 
        'risk_value': risk_data['risk_values']
    }
    
    for field_name, values in numeric_checks.items():
        non_numeric = []
        for item in values:
            if not isinstance(item['value'], (int, float)):
                non_numeric.append(item)
        
        if non_numeric:
            verification_results['risk_values_numeric'] = False
            verification_results['issues'].append(f"{field_name}字段包含非数值：{[item['raw_text'] for item in non_numeric]}")
            verification_results['data_logic_score'] -= 1
    
    verification_results['data_summary']['risk_numeric_fields'] = {
        'possibility_count': len(risk_data['possibility_values']),
        'consequence_count': len(risk_data['consequence_values']),
        'risk_value_count': len(risk_data['risk_values'])
    }
    
    # 检查风险等级
    invalid_levels = []
    for level_item in risk_data['risk_levels']:
        level = level_item['level']
        if level not in VALIDATION_RULES['risk_levels']:
            invalid_levels.append(level_item)
    
    if invalid_levels:
        verification_results['risk_levels_valid'] = False
        verification_results['issues'].append(f"风险等级不符合标准：{[item['level'] for item in invalid_levels]}")
        verification_results['data_logic_score'] -= 1
    
    verification_results['data_summary']['risk_levels'] = {
        'total': len(risk_data['risk_levels']),
        'valid': len(risk_data['risk_levels']) - len(invalid_levels),
        'invalid': len(invalid_levels)
    }
    
    # 检查数据完整性
    if not potential_results:
        verification_results['issues'].append("缺少电位测试结果数据")
        verification_results['data_logic_score'] -= 1
    
    if not any([risk_data['possibility_values'], risk_data['consequence_values'], risk_data['risk_values']]):
        verification_results['issues'].append("缺少风险评价数值数据")
        verification_results['data_logic_score'] -= 1
    
    if not risk_data['risk_levels']:
        verification_results['issues'].append("缺少风险等级数据")
        verification_results['data_logic_score'] -= 1
    
    verification_results['deduction_points'] = 5 - max(0, verification_results['data_logic_score'])
    verification_results['potential_results'] = potential_results
    verification_results['risk_data'] = risk_data
    
    return verification_results

# ========================== 核心：数据逻辑正确性负样本提示词模板 ==========================
def get_data_logic_prompt(verification_data, prompt_type, sample_idx):
    """生成数据逻辑检查的提示词模板"""
    prompt_templates = {
        1: f"""任务：为数据逻辑检查生成"电位测试结果超出范围"负样本（第{sample_idx}个变体）。
要求：
1. 修改电位测试结果，使其超出0.85V~1.2V标准范围（如0.7V或1.3V）；
2. 保持其他信息不变（如风险评价、时间等），仅让电位测试结果产生问题；
3. 确保修改后的电位值在文档中格式正确，仅让数据逻辑验证时出现错误；
4. 这是第{sample_idx}个变体，请生成与之前不同的具体电位数值。
原始信息：
---
电位测试结果：{[r['value'] for r in verification_data['potential_results']] if verification_data['potential_results'] else "无数据"}
---
输出：仅返回修改后的超出范围电位测试结果，格式为：
电位测试结果：[修改后的超出范围电位值]""",
        
        2: f"""任务：为数据逻辑检查生成"风险评价字段非数值"负样本（第{sample_idx}个变体）。
要求：
1. 修改风险评价中的可能性、后果值、风险值等字段，使其包含非数值内容（如文字、符号等）；
2. 保持其他信息不变（如电位测试、时间等），仅让风险评价数值字段产生问题；
3. 确保修改后的内容在文档中格式合理，仅让数据逻辑验证时出现错误；
4. 这是第{sample_idx}个变体，请生成与之前不同的具体非数值内容。
原始信息：
---
可能性值：{[r['value'] for r in verification_data['risk_data']['possibility_values']] if verification_data['risk_data']['possibility_values'] else "无数据"}
后果值：{[r['value'] for r in verification_data['risk_data']['consequence_values']] if verification_data['risk_data']['consequence_values'] else "无数据"}
风险值：{[r['value'] for r in verification_data['risk_data']['risk_values']] if verification_data['risk_data']['risk_values'] else "无数据"}
---
输出：仅返回修改后的非数值字段，格式为：
可能性：[修改后的非数值内容]
后果值：[修改后的非数值内容]
风险值：[修改后的非数值内容]""",
        
        3: f"""任务：为数据逻辑检查生成"风险等级不符合标准"负样本（第{sample_idx}个变体）。
要求：
1. 修改风险等级，使其不属于"低、中、较高、高"四个标准等级（如"中等"、"很高"等）；
2. 保持其他信息不变（如电位测试、数值等），仅让风险等级产生问题；
3. 确保修改后的等级在文档中格式合理，仅让数据逻辑验证时出现错误；
4. 这是第{sample_idx}个变体，请生成与之前不同的具体非标准等级。
原始信息：
---
风险等级：{[r['level'] for r in verification_data['risk_data']['risk_levels']] if verification_data['risk_data']['risk_levels'] else "无数据"}
---
输出：仅返回修改后的非标准风险等级，格式为：
风险等级：[修改后的非标准等级]""",
        
        4: f"""任务：为数据逻辑检查生成"关键数据缺失"负样本（第{sample_idx}个变体）。
要求：
1. 描述关键数据字段（电位测试结果、风险评价数值、风险等级等）缺失的情况；
2. 保持其他信息不变，仅让关键数据产生缺失问题；
3. 明确指出哪些关键数据缺失，描述要专业、准确；
4. 这是第{sample_idx}个变体，请生成与之前不同的具体缺失情况。
原始信息：
---
电位测试结果：{[r['value'] for r in verification_data['potential_results']] if verification_data['potential_results'] else "无数据"}
风险评价数据：{[r['value'] for r in verification_data['risk_data']['possibility_values']] if verification_data['risk_data']['possibility_values'] else "无数据"}
风险等级：{[r['level'] for r in verification_data['risk_data']['risk_levels']] if verification_data['risk_data']['risk_levels'] else "无数据"}
---
输出：仅返回缺失的关键数据描述，格式为：
缺失数据：[具体缺失的数据类型和描述]"""
    }
    return prompt_templates.get(prompt_type, "")

def generate_structured_samples_with_model(verification_data, samples_per_type=3):
    """使用模型生成训练样本，每个类型生成多个变体"""
    samples = []
    
    # 检查是否有足够的数据生成样本
    has_potential_data = len(verification_data['potential_results']) > 0
    has_risk_data = (len(verification_data['risk_data']['possibility_values']) > 0 or 
                    len(verification_data['risk_data']['consequence_values']) > 0 or 
                    len(verification_data['risk_data']['risk_values']) > 0)
    has_risk_levels = len(verification_data['risk_data']['risk_levels']) > 0
    
    if not (has_potential_data or has_risk_data or has_risk_levels):
        print("⚠️ 警告：未找到足够的数据生成数据逻辑检查样本")
        return samples
    
    # 生成电位测试结果超出范围样本（多个变体）
    # 需求：即使未提取到电位数据，也必须生成该类样本
    print(f"🔄 正在生成电位测试结果超出范围负样本（{samples_per_type}个变体）...")
    for i in range(samples_per_type):
        print(f"  → 生成第{i+1}/{samples_per_type}个电位样本...")
        prompt = get_data_logic_prompt(verification_data, 1, i+1)
        chat_message = [{"role": "user", "content": prompt}]
        response = get_answer(chat_message)
        
        if response.strip():
            try:
                # 解析模型返回的不一致信息
                lines = response.strip().split('\n')
                potential_info = ""
                
                for line in lines:
                    if "电位测试结果：" in line:
                        potential_info = line.replace("电位测试结果：", "").strip()
                
                if potential_info:
                    # 构造用户问题
                    user_content = f"电位测试结果是否符合标准范围（0.85V~1.2V）。\n\n电位测试结果：{potential_info}"
                    
                    # 构造助手回答（确保给出标准范围内的正确示例值）
                    in_range_values = [r['value'] for r in verification_data['potential_results']
                                       if VALIDATION_RULES['potential_range'][0] <= r['value'] <= VALIDATION_RULES['potential_range'][1]]
                    if in_range_values:
                        # 使用已提取且在标准范围内的值
                        correct_potentials = ', '.join(str(v) for v in in_range_values)
                    else:
                        # 若没有可用的在范围内的原始值，提供一个安全的标准值示例
                        correct_potentials = "1.00"
                    assistant_content = f"检测到数据逻辑问题：电位测试结果超出标准范围。\n正确的描述应该是：{correct_potentials}"
                    
                    samples.append({
                        "messages": [
                            {"role": "user", "content": user_content},
                            {"role": "assistant", "content": assistant_content}
                        ],
                        "images": []
                    })
                    print(f"    ✅ 电位样本{i+1}生成成功")
                else:
                    raise ValueError("无法解析模型返回的电位信息")
                    
            except Exception as e:
                print(f"    ⚠️ 电位样本{i+1}处理异常：{e}")
        else:
            print(f"    ⚠️ 电位样本{i+1}生成失败")
        
        time.sleep(1)  # API限流控制
    
    # 生成风险评价字段非数值样本（多个变体）
    # 需求：即使未提取到风险数值数据，也必须生成该类样本
    print(f"🔄 正在生成风险评价字段非数值负样本（{samples_per_type}个变体）...")
    for i in range(samples_per_type):
        print(f"  → 生成第{i+1}/{samples_per_type}个风险数值样本...")
        prompt = get_data_logic_prompt(verification_data, 2, i+1)
        chat_message = [{"role": "user", "content": prompt}]
        response = get_answer(chat_message)
        
        if response.strip():
            try:
                # 解析模型返回的不一致信息
                lines = response.strip().split('\n')
                possibility_info = ""
                consequence_info = ""
                risk_value_info = ""
                
                for line in lines:
                    if "可能性：" in line:
                        possibility_info = line.replace("可能性：", "").strip()
                    elif "后果值：" in line:
                        consequence_info = line.replace("后果值：", "").strip()
                    elif "风险值：" in line:
                        risk_value_info = line.replace("风险值：", "").strip()
                
                if possibility_info or consequence_info or risk_value_info:
                    # 构造用户问题
                    user_content = f"风险评价中的数值字段是否为有效数值。\n\n可能性：{possibility_info}\n后果值：{consequence_info}\n风险值：{risk_value_info}"
                    
                    # 构造助手回答
                    original_possibility = [str(r['value']) for r in verification_data['risk_data']['possibility_values']]
                    original_consequence = [str(r['value']) for r in verification_data['risk_data']['consequence_values']]
                    original_risk_value = [str(r['value']) for r in verification_data['risk_data']['risk_values']]
                    correct_description = f"可能性：{', '.join(original_possibility) if original_possibility else '无数据'}，后果值：{', '.join(original_consequence) if original_consequence else '无数据'}，风险值：{', '.join(original_risk_value) if original_risk_value else '无数据'}"
                    assistant_content = f"检测到数据逻辑问题：风险评价字段包含非数值内容。\n正确的描述应该是：{correct_description}"
                    
                    samples.append({
                        "messages": [
                            {"role": "user", "content": user_content},
                            {"role": "assistant", "content": assistant_content}
                        ],
                        "images": []
                    })
                    print(f"    ✅ 风险数值样本{i+1}生成成功")
                else:
                    raise ValueError("无法解析模型返回的风险数值信息")
                    
            except Exception as e:
                print(f"    ⚠️ 风险数值样本{i+1}处理异常：{e}")
        else:
            print(f"    ⚠️ 风险数值样本{i+1}生成失败")
        
        time.sleep(1)  # API限流控制
    
    # 生成风险等级不符合标准样本（多个变体）
    # 需求：即使未提取到风险等级数据，也必须生成该类样本
    print(f"🔄 正在生成风险等级不符合标准负样本（{samples_per_type}个变体）...")
    for i in range(samples_per_type):
        print(f"  → 生成第{i+1}/{samples_per_type}个风险等级样本...")
        prompt = get_data_logic_prompt(verification_data, 3, i+1)
        chat_message = [{"role": "user", "content": prompt}]
        response = get_answer(chat_message)
        
        if response.strip():
            try:
                # 解析模型返回的不一致信息
                lines = response.strip().split('\n')
                risk_level_info = ""
                
                for line in lines:
                    if "风险等级：" in line:
                        risk_level_info = line.replace("风险等级：", "").strip()
                
                if risk_level_info:
                    # 构造用户问题
                    user_content = f"风险等级是否符合标准（低、中、较高、高）。\n\n风险等级：{risk_level_info}"
                    
                    # 构造助手回答
                    original_levels = [r['level'] for r in verification_data['risk_data']['risk_levels']]
                    correct_levels = ', '.join(original_levels) if original_levels else "无数据"
                    assistant_content = f"检测到数据逻辑问题：风险等级不符合标准。\n正确的描述应该是：{correct_levels}"
                    
                    samples.append({
                        "messages": [
                            {"role": "user", "content": user_content},
                            {"role": "assistant", "content": assistant_content}
                        ],
                        "images": []
                    })
                    print(f"    ✅ 风险等级样本{i+1}生成成功")
                else:
                    raise ValueError("无法解析模型返回的风险等级信息")
                    
            except Exception as e:
                print(f"    ⚠️ 风险等级样本{i+1}处理异常：{e}")
        else:
            print(f"    ⚠️ 风险等级样本{i+1}生成失败")
        
        time.sleep(1)  # API限流控制
    
    # 按需求：不再生成“关键数据缺失”类负样本
    
    return samples

# ========================== 主程序 ==========================
if __name__ == '__main__':
    document_paths = [INPUT_DOCUMENT_PATH]
    
    try:
        file_name_for_prompt = os.path.basename(INPUT_DOCUMENT_PATH)
        print(f"准备处理文档: {file_name_for_prompt}")
    except:
        file_name_for_prompt = "高后果区管控方案"
        print(f"准备处理文档: {INPUT_DOCUMENT_PATH}")
    
    all_generated_data = []
    for doc_path in document_paths:
        print(f"\n{'='*80}\n正在处理文件: {doc_path}\n{'='*80}")
        
        # 提取Word文档文本
        full_text = extract_text_from_docx(doc_path)
        if not full_text:
            print("Word文档文本提取失败")
            continue
        
        print(f"成功提取文本，长度: {len(full_text)}字符")
        
        # 提取电位测试结果
        potential_results = extract_potential_test_results(full_text)
        print(f"提取到电位测试结果: {len(potential_results)}个")
        
        # 提取风险评价数据
        risk_data = extract_risk_assessment_data(full_text)
        total_risk_data = (len(risk_data['possibility_values']) + 
                          len(risk_data['consequence_values']) + 
                          len(risk_data['risk_values']) + 
                          len(risk_data['risk_levels']))
        print(f"提取到风险评价数据: {total_risk_data}个")
        
        # 验证数据逻辑
        verification = verify_data_logic(potential_results, risk_data)
        
        print(f"数据逻辑评分: {verification['data_logic_score']}/5")
        if verification['issues']:
            print(f"发现问题: {verification['issues']}")
        
        # 使用模型生成负样本
        print(f"\n🤖 开始使用模型生成数据逻辑检查负样本...")
        print(f"📝 配置：每个负样本类型将生成 3 个变体")
        samples = generate_structured_samples_with_model(verification, 3)
        all_generated_data.extend(samples)
    
    # 写入输出文件
    if all_generated_data:
        output_dir = os.path.dirname(OUTPUT_JSON_PATH)
        if not os.path.exists(output_dir):
            try:
                os.makedirs(output_dir)
                print(f"已创建输出目录: {output_dir}")
            except OSError as e:
                print(f"错误: 创建输出目录失败: {e}")
                all_generated_data = []
        
        if all_generated_data:
            try:
                with open(OUTPUT_JSON_PATH, "w", encoding="utf-8") as f:
                    json.dump(all_generated_data, f, ensure_ascii=False, indent=4)
                print(f"\n{'='*80}")
                print(f"✅ 处理完成！数据保存至：{OUTPUT_JSON_PATH}")
                print(f"共生成 {len(all_generated_data)} 个对话样本（每个包含负样本user + 正样本assistant）")
                
                # 统计各类型样本数量
                potential_count = sum(1 for s in all_generated_data if "电位测试结果超出标准范围" in s["messages"][1]["content"])
                risk_numeric_count = sum(1 for s in all_generated_data if "风险评价字段包含非数值内容" in s["messages"][1]["content"])
                risk_level_count = sum(1 for s in all_generated_data if "风险等级不符合标准" in s["messages"][1]["content"])
                print(f"样本分布：电位超出范围 {potential_count} 个，风险字段非数值 {risk_numeric_count} 个，风险等级不符 {risk_level_count} 个")
                print("="*80 + "\n")
            except Exception as e:
                print(f"错误: 写入JSON失败: {e}")
    else:
        print("\n\n处理结束，无生成数据。请检查Word文档或API配置。")
